# 加密工具类
import os
import hashlib
import secrets
import string
import base64
from Crypto.Cipher import AES
from Crypto.PublicKey import RSA
from Crypto import Random
from Crypto.Util.Padding import pad, unpad
from typing import Dict, Tuple, Union, Optional


class CryptoUtils:
    """加密工具类，提供各种加密算法的密钥生成和加密/解密功能"""
    
    @staticmethod
    def generate_md5_hash(data: str) -> str:
        """
        生成MD5哈希值
        
        Args:
            data: 要哈希的数据
            
        Returns:
            MD5哈希字符串
        """
        return hashlib.md5(data.encode('utf-8')).hexdigest()
    
    @staticmethod
    def generate_sha1_hash(data: str) -> str:
        """
        生成SHA1哈希值
        
        Args:
            data: 要哈希的数据
            
        Returns:
            SHA1哈希字符串
        """
        return hashlib.sha1(data.encode('utf-8')).hexdigest()
    
    @staticmethod
    def generate_sha256_hash(data: str) -> str:
        """
        生成SHA256哈希值
        
        Args:
            data: 要哈希的数据
            
        Returns:
            SHA256哈希字符串
        """
        return hashlib.sha256(data.encode('utf-8')).hexdigest()
    
    @staticmethod
    def generate_sha512_hash(data: str) -> str:
        """
        生成SHA512哈希值
        
        Args:
            data: 要哈希的数据
            
        Returns:
            SHA512哈希字符串
        """
        return hashlib.sha512(data.encode('utf-8')).hexdigest()
    
    @staticmethod
    def generate_random_secret(length: int = 32) -> str:
        """
        生成随机密钥
        
        Args:
            length: 密钥长度，默认32位
            
        Returns:
            随机密钥字符串
        """
        alphabet = string.ascii_letters + string.digits + "!@#$%^&*"
        return ''.join(secrets.choice(alphabet) for _ in range(length))
    
    @staticmethod
    def generate_hex_secret(length: int = 32) -> str:
        """
        生成十六进制随机密钥
        
        Args:
            length: 密钥长度，默认32位
            
        Returns:
            十六进制随机密钥字符串
        """
        return secrets.token_hex(length)
    
    @staticmethod
    def generate_base64_secret(length: int = 32) -> str:
        """
        生成Base64编码的随机密钥
        
        Args:
            length: 密钥长度，默认32位
            
        Returns:
            Base64编码的随机密钥字符串
        """
        return base64.b64encode(secrets.token_bytes(length)).decode('utf-8')
    
    @staticmethod
    def generate_rsa_key_pair(key_size: int = 2048) -> Tuple[str, str]:
        """
        生成RSA密钥对
        
        Args:
            key_size: 密钥长度，默认2048位
            
        Returns:
            (私钥, 公钥) 元组
        """
        key = RSA.generate(key_size)
        private_key = key.export_key().decode('utf-8')
        public_key = key.publickey().export_key().decode('utf-8')
        return private_key, public_key
    
    @staticmethod
    def generate_aes_key(key_size: int = 256) -> str:
        """
        生成AES密钥
        
        Args:
            key_size: 密钥长度，支持128、192、256位，默认256位
            
        Returns:
            AES密钥字符串
        """
        if key_size not in [128, 192, 256]:
            raise ValueError("AES密钥长度必须是128、192或256位")
        
        key_bytes = key_size // 8
        return secrets.token_hex(key_bytes)
    
    @staticmethod
    def generate_des_key() -> str:
        """
        生成DES密钥
        
        Returns:
            DES密钥字符串
        """
        return secrets.token_hex(8)  # DES密钥是8字节
    
    @staticmethod
    def aes_encrypt(plaintext: str, key: str) -> str:
        """
        使用AES加密数据
        
        Args:
            plaintext: 要加密的明文
            key: AES密钥，十六进制字符串
            
        Returns:
            Base64编码的加密数据
        """
        key_bytes = bytes.fromhex(key)
        iv = Random.new().read(AES.block_size)
        cipher = AES.new(key_bytes, AES.MODE_CBC, iv)
        padded_text = pad(plaintext.encode('utf-8'), AES.block_size)
        encrypted_text = cipher.encrypt(padded_text)
        return base64.b64encode(iv + encrypted_text).decode('utf-8')
    
    @staticmethod
    def aes_decrypt(ciphertext: str, key: str) -> str:
        """
        使用AES解密数据
        
        Args:
            ciphertext: Base64编码的加密数据
            key: AES密钥，十六进制字符串
            
        Returns:
            解密后的明文
        """
        key_bytes = bytes.fromhex(key)
        encrypted_data = base64.b64decode(ciphertext)
        iv = encrypted_data[:AES.block_size]
        cipher = AES.new(key_bytes, AES.MODE_CBC, iv)
        decrypted_text = unpad(cipher.decrypt(encrypted_data[AES.block_size:]), AES.block_size)
        return decrypted_text.decode('utf-8')
    
    @staticmethod
    def rsa_encrypt(plaintext: str, public_key: str) -> str:
        """
        使用RSA公钥加密数据
        
        Args:
            plaintext: 要加密的明文
            public_key: RSA公钥
            
        Returns:
            Base64编码的加密数据
        """
        key = RSA.import_key(public_key)
        cipher_text = key.encrypt(plaintext.encode('utf-8'), 32)[0]
        return base64.b64encode(cipher_text).decode('utf-8')
    
    @staticmethod
    def rsa_decrypt(ciphertext: str, private_key: str) -> str:
        """
        使用RSA私钥解密数据
        
        Args:
            ciphertext: Base64编码的加密数据
            private_key: RSA私钥
            
        Returns:
            解密后的明文
        """
        key = RSA.import_key(private_key)
        encrypted_data = base64.b64decode(ciphertext)
        decrypted_text = key.decrypt(encrypted_data, 32)
        return decrypted_text.decode('utf-8')
    
    @staticmethod
    def generate_hmac_hash(data: str, key: str, algorithm: str = 'sha256') -> str:
        """
        生成HMAC哈希值
        
        Args:
            data: 要哈希的数据
            key: 密钥
            algorithm: 哈希算法，支持md5, sha1, sha224, sha256, sha384, sha512
            
        Returns:
            HMAC哈希字符串
        """
        if algorithm not in hashlib.algorithms_available:
            raise ValueError(f"不支持的哈希算法: {algorithm}")
        
        hmac = hashlib.pbkdf2_hmac(
            algorithm,
            data.encode('utf-8'),
            key.encode('utf-8'),
            100000  # 迭代次数
        )
        return base64.b64encode(hmac).decode('utf-8')
    
    @staticmethod
    def generate_pbkdf2_hash(password: str, salt: Optional[str] = None, iterations: int = 100000, 
                            key_length: int = 32, algorithm: str = 'sha256') -> Tuple[str, str]:
        """
        使用PBKDF2算法生成密码哈希
        
        Args:
            password: 要哈希的密码
            salt: 盐值，如果不提供则随机生成
            iterations: 迭代次数
            key_length: 密钥长度
            algorithm: 哈希算法
            
        Returns:
            (哈希值, 盐值) 元组
        """
        if algorithm not in hashlib.algorithms_available:
            raise ValueError(f"不支持的哈希算法: {algorithm}")
        
        if salt is None:
            salt = secrets.token_hex(16)
        
        hash_value = hashlib.pbkdf2_hmac(
            algorithm,
            password.encode('utf-8'),
            salt.encode('utf-8'),
            iterations,
            key_length
        )
        
        return base64.b64encode(hash_value).decode('utf-8'), salt
    
    @staticmethod
    def generate_jwt_secret(algorithm: str = 'random', length: int = 32) -> str:
        """
        生成JWT密钥
        
        Args:
            algorithm: 生成算法，支持random, hex, base64, md5, sha1, sha256, sha512
            length: 密钥长度，适用于random, hex, base64算法
            
        Returns:
            JWT密钥字符串
        """
        if algorithm == 'random':
            return CryptoUtils.generate_random_secret(length)
        elif algorithm == 'hex':
            return CryptoUtils.generate_hex_secret(length)
        elif algorithm == 'base64':
            return CryptoUtils.generate_base64_secret(length)
        elif algorithm == 'md5':
            return CryptoUtils.generate_md5_hash(secrets.token_hex(16))
        elif algorithm == 'sha1':
            return CryptoUtils.generate_sha1_hash(secrets.token_hex(16))
        elif algorithm == 'sha256':
            return CryptoUtils.generate_sha256_hash(secrets.token_hex(16))
        elif algorithm == 'sha512':
            return CryptoUtils.generate_sha512_hash(secrets.token_hex(16))
        else:
            raise ValueError(f"不支持的算法: {algorithm}")


# 示例用法
if __name__ == "__main__":
    # 生成各种密钥示例
    print("MD5哈希:", CryptoUtils.generate_md5_hash("test"))
    print("SHA1哈希:", CryptoUtils.generate_sha1_hash("test"))
    print("SHA256哈希:", CryptoUtils.generate_sha256_hash("test"))
    print("SHA512哈希:", CryptoUtils.generate_sha512_hash("test"))
    
    print("\n随机密钥:", CryptoUtils.generate_random_secret(32))
    print("十六进制密钥:", CryptoUtils.generate_hex_secret(32))
    print("Base64密钥:", CryptoUtils.generate_base64_secret(32))
    
    # 生成AES密钥
    aes_key = CryptoUtils.generate_aes_key(256)
    print("\nAES密钥:", aes_key)
    
    # AES加密/解密示例
    plaintext = "这是一段需要加密的文本"
    encrypted = CryptoUtils.aes_encrypt(plaintext, aes_key)
    print("AES加密:", encrypted)
    decrypted = CryptoUtils.aes_decrypt(encrypted, aes_key)
    print("AES解密:", decrypted)
    
    # 生成RSA密钥对
    private_key, public_key = CryptoUtils.generate_rsa_key_pair(2048)
    print("\nRSA私钥:", private_key[:50] + "...")
    print("RSA公钥:", public_key[:50] + "...")
    
    # RSA加密/解密示例
    rsa_encrypted = CryptoUtils.rsa_encrypt("RSA加密测试", public_key)
    print("RSA加密:", rsa_encrypted)
    rsa_decrypted = CryptoUtils.rsa_decrypt(rsa_encrypted, private_key)
    print("RSA解密:", rsa_decrypted)
    
    # 生成JWT密钥
    print("\nJWT密钥(random):", CryptoUtils.generate_jwt_secret('random', 32))
    print("JWT密钥(hex):", CryptoUtils.generate_jwt_secret('hex', 32))
    print("JWT密钥(sha256):", CryptoUtils.generate_jwt_secret('sha256'))